from dataclasses import asdict, dataclass, field
from enum import Enum
from typing import Dict, List, Type, Union

import orjson as json

from tcp_comm.common import *
from tcp_comm.utils import restore_from_dict


class MsgType(Enum):
    MsgType_None = 0
    MsgType_ObjectInit = 1
    MsgType_ObjectInfo = 2
    MsgType_TaskStatus = 3
    MsgType_SetMsgFrequency = 4
    MsgType_CreateGroup = 5
    MsgType_SetCombatMode = 6
    MsgType_DefendAreaInfo = 7
    MsgType_SetGroupTarget = 8
    MsgType_ManeuveringAction = 9
    MsgType_FireAssault = 10
    MsgType_RandomTargetGeneration = 11
    MsgType_ObjectPerceptionInfo = 12
    MsgType_SimulationResult = 13
    MsgType_EntityChangedInfo = 14
    MsgType_EventInfo = 15
    MsgType_DynamicObstacles = 16
    MsgType_UavInfo = 17
    MsgType_SimulationSpeed = 18
    MsgType_StatusVerify = 19
    MsgType_SetAmmoCount = 20
    MsgType_NoFlyZoneInfo = 21
    MsgType_AllocateThreat = 22
    MsgType_DestoryObject = 23
    MsgType_SetGroupTargets = 24
    MsgType_SetGroupFormation = 25


class EventType(Enum):
    EventType_Fire = 0
    EventType_Hit = 1
    EventType_Collision = 2


class MsgRegistry:
    _registry: Dict[MsgType, Type["MessageBase"]] = {}
    _event_registry: Dict[EventType, Type["EventInfoMsg"]] = {}

    @classmethod
    def register(cls, msg_type: MsgType, event_type: EventType = None):
        def decorator(message_cls):
            if event_type is not None:
                cls._event_registry[event_type] = message_cls
            else:
                cls._registry[msg_type] = message_cls

            return message_cls

        return decorator

    @classmethod
    def get_msg_class(cls, msg_type: MsgType, event_type: EventType = None):
        if event_type is not None:
            if event_type not in cls._event_registry:
                raise ValueError(f"Undefined event type: {event_type}")
            return cls._event_registry[event_type]
        elif msg_type not in cls._registry:
            raise ValueError(f"Undefined message type: {msg_type}")

        return cls._registry[msg_type]


@dataclass
class MessageBase:
    timestamps: float
    """Simulation time when the message is generated"""

    msgType: MsgType = field(init=False)
    """Message type"""

    def to_dict(self):
        """Convert message to json serializable dict

        Note: Only `fields` will be kept, attributes in __post_init__ will be lost
        """
        d = asdict(self)
        for key, value in d.items():
            if isinstance(value, Enum):
                d[key] = value.value
        return d

    @classmethod
    def from_json(cls, json_data: Union[str, bytes, dict]) -> "MessageBase":
        if isinstance(json_data, bytes):
            json_str = json_data.decode("utf-8")
        elif isinstance(json_data, str):
            json_str = json_data
        else:
            json_str = ""

        d = json.loads(json_str) if json_str else json_data
        if not isinstance(d, dict):
            raise ValueError("Invalid json data!\n", json_data)

        msg_type = MsgType(d["msgType"])
        eventType = EventType(d["eventType"]) if "eventType" in d else None
        msg_cls = MsgRegistry.get_msg_class(msg_type, eventType)
        if cls is not MessageBase and not msg_cls is cls:
            raise ValueError(f"Expect {cls} but got {msg_cls}!")

        message = restore_from_dict(msg_cls, d)
        return message


@MsgRegistry.register(MsgType.MsgType_EventInfo)
@dataclass
class EventInfoMsg(MessageBase):
    eventType: EventType = field(init=False)
    """Event type"""


@MsgRegistry.register(MsgType.MsgType_ObjectInit)
@dataclass
class ObjectInitActionMsg(MessageBase):
    objectName: str
    objectType: ObjectType
    objectCamp: int
    objectPos: Position3D
    objectOrientation: RotationAngles

    def __post_init__(self):
        self.msgType = MsgType.MsgType_ObjectInit


@MsgRegistry.register(MsgType.MsgType_ObjectInfo)
@dataclass
class ObjectInfoMsg(MessageBase):
    objectInfo: ObjectInfo

    def __post_init__(self):
        self.msgType = MsgType.MsgType_ObjectInfo


@MsgRegistry.register(MsgType.MsgType_TaskStatus)
@dataclass
class TaskStatusInfoMsg(MessageBase):
    taskStatus: int
    """0: Server ready, 1: Start simulation, 2: request result, 3: request new episode, 4: shutdown"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_TaskStatus


@MsgRegistry.register(MsgType.MsgType_SetMsgFrequency)
@dataclass
class SetMsgFrequencyMsg(MessageBase):
    frequency: int
    """milliseconds of message frequency"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetMsgFrequency


@MsgRegistry.register(MsgType.MsgType_CreateGroup)
@dataclass
class CreateGroupInfoMsg(MessageBase):
    objectNameList: str
    """A string list of objectName, name is separated by '#'."""

    groupID: int

    def __post_init__(self):
        self.msgType = MsgType.MsgType_CreateGroup

        if self.objectNameList.startswith("["):
            names = json.loads(self.objectNameList)
            self.objectNameList = "#".join(names)


@MsgRegistry.register(MsgType.MsgType_SetCombatMode)
@dataclass
class SetCombatModeMsg(MessageBase):
    groupID: int
    mode: int

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetCombatMode


@MsgRegistry.register(MsgType.MsgType_DefendAreaInfo)
@dataclass
class DefendAreaInfoMsg(MessageBase):
    objectName: str
    defendAreaInfo: DefendAreaInfo

    def __post_init__(self):
        self.msgType = MsgType.MsgType_DefendAreaInfo


@MsgRegistry.register(MsgType.MsgType_SetGroupTarget)
@dataclass
class SetGroupTargetInfoMsg(MessageBase):
    groupID: int
    targetPos: Position3D

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetGroupTarget


@MsgRegistry.register(MsgType.MsgType_ManeuveringAction)
@dataclass
class ManeuveringActionMsg(MessageBase):
    objectName: str
    """The object name of the object that needs to perform the maneuvering action"""

    path: str
    """Json serialized list of ManeuveringInfo"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_ManeuveringAction
        self.path_points = [ManeuveringInfo(**x) for x in json.loads(self.path)]


@MsgRegistry.register(MsgType.MsgType_FireAssault)
@dataclass
class FireAssaultActionMsg(MessageBase):
    startPos: Position3D
    """Task start position"""
    aimAreaTopLeftPos: Position3D
    """The top left position of the aim area"""
    aimAreaBottomRightPos: Position3D
    """The bottom right position of the aim area"""
    pathPosList: str
    """Json serialized list of Position3D"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_FireAssault
        self.path_points = [Position3D(**x) for x in json.loads(self.pathPosList)]


@MsgRegistry.register(MsgType.MsgType_RandomTargetGeneration)
@dataclass
class RandomTargetGenerationMsg(MessageBase):
    targetType: ObjectType
    targetNum: int
    areaTopLeftPoint: Position3D
    areaBottomRightPoint: Position3D
    successPossibility: float = 1.0

    def __post_init__(self):
        self.msgType = MsgType.MsgType_RandomTargetGeneration


@MsgRegistry.register(MsgType.MsgType_ObjectPerceptionInfo)
@dataclass
class ObjectPerceptionInfoMsg(MessageBase):
    objectName: str
    entityNum: int
    entityDiscovered: str
    """A string list of objectName, name is separated by '#'."""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_ObjectPerceptionInfo

        if self.entityDiscovered.startswith("["):
            self.entities: List[str] = json.loads(self.entityDiscovered)
        else:
            self.entities: List[str] = self.entityDiscovered.split("#")


@MsgRegistry.register(MsgType.MsgType_SimulationResult)
@dataclass
class SimulationResultMsg(MessageBase):
    duration: float
    """Simulation duration in seconds"""
    isTaskCompleted: bool
    """Whether the task is completed"""
    redUnitsRemains: int
    """Total number of red units that remain"""
    redUnitsRemainsInfo: str
    """Json serialized list of [(red unit type, red unit number)]"""
    blueUnitsRemains: int
    """Total number of blue units that remain"""
    blueUnitsRemainsInfo: str
    """Json serialized list of [(blue unit type, blue unit number)]"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SimulationResult


@MsgRegistry.register(MsgType.MsgType_EntityChangedInfo)
@dataclass
class EntityChangedInfoMsg(MessageBase):
    camp: int
    changedNum: int
    entityInfoVec: str
    """Json serialized list of ObjectInfo"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_EntityChangedInfo
        self.entities = [ObjectInfo(**x) for x in json.loads(self.entityInfoVec)]


@MsgRegistry.register(MsgType.MsgType_EventInfo, EventType.EventType_Collision)
@dataclass
class CollisionInfoMsg(EventInfoMsg):
    objectName: str

    def __post_init__(self):
        self.eventType = EventType.EventType_Collision
        self.msgType = MsgType.MsgType_EventInfo


@MsgRegistry.register(MsgType.MsgType_EventInfo, EventType.EventType_Hit)
@dataclass
class HitInfoMsg(EventInfoMsg):
    shooterName: str
    targetName: str
    damage: float

    def __post_init__(self):
        self.eventType = EventType.EventType_Hit
        self.msgType = MsgType.MsgType_EventInfo


@MsgRegistry.register(MsgType.MsgType_DynamicObstacles)
@dataclass
class DynamicObstaclesInfoMsg(MessageBase):
    objectName: str
    """The object that perceives the dynamic obstacles"""
    obstacles: str
    """Json serialized list of Vector3D (In grid coordinate)."""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_DynamicObstacles

        obstacles = json.loads(self.obstacles)
        self.obstacle_pos = [Vector3D(**x) for x in obstacles]


@MsgRegistry.register(MsgType.MsgType_UavInfo)
@dataclass
class UAVInfoMsg(MessageBase):
    objectInfo: ObjectInfo
    obstacles: str
    """List of dynamic obstacles occupancies in global coordinate"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_UavInfo

        obstacles = json.loads(self.obstacles)
        self.obstacles_pos = [(Vector3D(**bl), Vector3D(**tr)) for bl, tr in obstacles]


@MsgRegistry.register(MsgType.MsgType_SimulationSpeed)
@dataclass
class SimulationSpeedInfoMsg(MessageBase):
    speed: int
    """Toggle the simulation speed to x times (maximum 32x)"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SimulationSpeed


@MsgRegistry.register(MsgType.MsgType_StatusVerify)
@dataclass
class StatusVerifyInfoMsg(MessageBase):
    statusVerify: int
    """0-ObjectInit Success, 1-Create Group Success, 2-Set Group Target Success"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_StatusVerify


@MsgRegistry.register(MsgType.MsgType_SetAmmoCount)
@dataclass
class SetAmmoCountInfoMsg(MessageBase):
    objectName: str
    ammoCount: str
    """Json serialized list of AmmoCountInfo"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetAmmoCount


@MsgRegistry.register(MsgType.MsgType_NoFlyZoneInfo)
@dataclass
class SetNoFlyZoneInfoMsg(MessageBase):
    type: int
    """Type of no-fly zone, 1: cube with red color"""
    leftBottom: Position3D
    """Left bottom position of the no-fly zone"""
    rightTop: Position3D
    """Right top position of the no-fly zone"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_NoFlyZoneInfo


@MsgRegistry.register(MsgType.MsgType_AllocateThreat)
@dataclass
class AllocateThreatInfoMsg(MessageBase):
    """Send `objectName` drone to `targetPos` threat"""

    objectName: str
    targetPos: Position3D
    targetDir: float = 100  # for compatability

    def __post_init__(self):
        self.msgType = MsgType.MsgType_AllocateThreat


@MsgRegistry.register(MsgType.MsgType_DestoryObject)
@dataclass
class DestoryObjectInfoMsg(MessageBase):
    objectNameList: str
    """Destory object name str list"""

    def __post_init__(self):
        self.msgType = MsgType.MsgType_DestoryObject


@MsgRegistry.register(MsgType.MsgType_SetGroupTargets)
@dataclass
class SetGroupTargetsInfoMsg(MessageBase):
    groupID: int
    targetPos: List[Position3D]

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetGroupTargets


@MsgRegistry.register(MsgType.MsgType_SetGroupFormation)
@dataclass
class SetGroupFormationMsg(MessageBase):
    groupID: int
    formationID: int

    def __post_init__(self):
        self.msgType = MsgType.MsgType_SetGroupFormation
